import cryptography.hazmat.primitives.cmac as cmac
import cryptography.hazmat.primitives.ciphers.algorithms as alg

key_nrf53 = [
    0xc5, 0xa8, 0x08, 0xeb, 0xe3, 0x1e, 0xa5, 0xb4,
    0xe9, 0x44, 0x1a, 0x76, 0x45, 0x58, 0xd8, 0x8b,
    0x40, 0x26, 0x33, 0xa8, 0xcd, 0x2d, 0x51, 0x67,
    0x8d, 0xef, 0x00, 0x24, 0x30, 0x52, 0xd7, 0x3d]

label_53 = [
    0xff, 0xff, 0xff, 0xff, # psa_crypto puts the partition id first.
    0xe2, 0x75, 0xab, 0x25, 0x00, 0x3b, 0x15, 0xe1,
    0xa1, 0x98, 0x78, 0x5b, 0x4e, 0x57, 0x43, 0x9a]

key_nrf91 = [
    0x19, 0x9a, 0xe3, 0xc7, 0x9d, 0xd0, 0x16, 0x8c,
    0x3e, 0xee, 0xa8, 0x46, 0xea, 0x4e, 0xdc, 0x6e]

label_91 = [
    0xff, 0xff, 0xff, 0xff, # psa_crypto puts the partition id first.
    0x29, 0xd6, 0xa6, 0x8d, 0x3f, 0xfb, 0x39, 0x1f,
    0x03, 0x6d, 0xd5, 0x3b, 0xe3, 0x54, 0xe4, 0x02,
    0x9a, 0x69, 0x5a]

context_53 = []

cmac_53 = cmac.CMAC(alg.AES(bytes(key_nrf53)), backend=None)
cmac_53.update(b'\1' + bytes(label_53) + b'\0' + bytes(context_53) + b'\1\0')
result = cmac_53.finalize()

cmac_53_2 = cmac.CMAC(alg.AES(bytes(key_nrf53)), backend=None)
cmac_53_2.update(b'\2' + bytes(label_53) + b'\0' + bytes(context_53) + b'\1\0')
result += cmac_53_2.finalize()

print([hex(c) for c in result])

context_91 = []

cmac_91_2 = cmac.CMAC(alg.AES(bytes(key_nrf91)), backend=None)
cmac_91_2.update(b'\1' + bytes(label_91) + b'\0' + bytes(context_91) + b'\x80')
result = cmac_91_2.finalize()

print([hex(c) for c in result])
